import argparse
from transformers import BertModel
import torch
from torch.utils.data import Dataset
from transformers import BertTokenizer
import numpy as np
from models.BERT_BASE import BERT
from utils.data_utils import Tokenizer, process_string
import torch.nn as nn
import sys


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--dropout", type=float, default=0.1)
    parser.add_argument("--threshold",type=float,default=0.5)
    parser.add_argument("--num_classes", type=int, default=1)
    parser.add_argument("--max_seq_len", type=int, default=512)
    parser.add_argument("--pretrained_bert_name", type=str, default='bert-base-uncased')
    parser.add_argument("--state_dict_path",type=str,default="10:34:24_val_acc_96.72.pth")
    args = parser.parse_args()
    args.state_dict_path = "checkout/state_dict/" + args.state_dict_path
    args.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f"> loading pretrained model: {args.state_dict_path}")
    tokenizer = Tokenizer(args.max_seq_len, args.pretrained_bert_name)
    bert = BertModel.from_pretrained(args.pretrained_bert_name)
    model = BERT(bert, args).to(args.device)
    model.load_state_dict(torch.load(args.state_dict_path))
    model.eval()
    torch.autograd.set_grad_enabled(False)
    print("input your sentence: ")
    with torch.no_grad():
        while 1:
            # a = "the design and atmosphere is just as good."
            a = sys.stdin.readline().strip()
            if a == 'exit':
                break
            a = process_string(a)
            print(f"Part of the string after processing as input:\n"
                  f"{a[:100] if len(a)>100 else a}")
            token_list = tokenizer.text_to_ids(a)
            attention_mask = torch.tensor(
                [1 if x != 0 else 0 for x in token_list]).view(1,-1).to(args.device)
            inputs = torch.tensor(token_list).view(1,-1).to(args.device)
            outputs = model(input_ids=inputs,attention_mask=attention_mask)
            print(f"Positive probability is : {outputs.item()}")
            pred = outputs.item() > args.threshold
            if pred:
                print("Positive\n")
            else:
                print("Negative\n")


if __name__ == "__main__":
    main()